"""
视频标注中间件模块

该模块提供视频对象识别和关键词搜索功能
"""

from dataclasses import dataclass
from typing import List, Callable, Optional
import cv2
from pathlib import Path
from ultralytics import YOLO


@dataclass
class CMatchVideo:
    """
    关键词搜索匹配结果数据类
    
    用于存储在结果文件中搜索关键词后的匹配信息
    """
    m_videoName: str              # 视频文件名，如 "example.mp4"
    m_videoPath: str              # 视频完整路径，如 "videos/example.mp4"
    m_resultFile: str             # 结果文件名，如 "example.txt"
    m_matchedLines: List[dict]    # 匹配行列表，每个元素包含 'line_number' 和 'content'
    m_matchCnt: int              # 匹配次数


class VideoAnnotationMiddleware:
    """
    视频标注中间件核心类
    
    负责协调YOLO模型加载、视频处理和关键词搜索功能
    """
    
    def __init__(self, videos_dir: str = "videos", results_dir: str = "results", model_name: str = "yolo11n.pt"):
        """
        初始化VideoAnnotationMiddleware
        
        参数:
            videos_dir: 视频文件夹路径，默认 "videos"
            results_dir: 结果文件夹路径，默认 "results"
            model_name: YOLO模型名称，默认 "yolo11n.pt"
        """
        self.m_videosDir = videos_dir
        self.m_resultsDir = results_dir
        self.m_modelName = model_name
        self._model = None  # YOLO模型实例（私有变量）
    
    def loadModel(self) -> None:
        """
        加载YOLO模型到内存
        
        功能:
            - 使用ultralytics库加载指定的YOLO模型
            - 如果模型文件不存在，自动下载
            - 验证模型加载成功
        
        异常:
            RuntimeError: 模型加载失败
        """
        try:
            # 使用ultralytics.YOLO类加载模型
            # 如果模型文件不存在，ultralytics会自动从官方源下载
            print(f"正在加载YOLO模型: {self.m_modelName}")
            self._model = YOLO(self.m_modelName)
            
            # 验证模型加载成功
            if self._model is None:
                raise RuntimeError(f"YOLO模型加载失败: {self.m_modelName}")
            
            print(f"YOLO模型加载成功: {self.m_modelName}")
            
        except Exception as e:
            # 添加模型加载失败的异常处理
            error_msg = f"加载YOLO模型时出错 ({self.m_modelName}): {str(e)}"
            print(error_msg)
            raise RuntimeError(error_msg) from e
    
    def processVideo(self, video_path: str, callback: Optional[Callable[[int, int], None]] = None) -> bool:
        """
        处理单个视频文件
        
        参数:
            video_path: 视频文件的完整路径
            callback: 可选的回调函数，签名为 callback(current_frame, total_frames)
        
        返回值:
            bool: 处理成功返回True，失败返回False
        
        功能:
            - 使用OpenCV读取视频
            - 逐帧进行YOLO检测
            - 将结果写入txt文件
            - 每处理一帧调用callback（如果提供）
        """
        cap = None
        try:
            # 子任务 5.1: 使用cv2.VideoCapture打开视频文件
            cap = cv2.VideoCapture(video_path)
            
            # 检查视频是否成功打开
            if not cap.isOpened():
                raise ValueError(f"无法打开视频文件: {video_path}")
            
            # 获取视频总帧数用于进度计算
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            # 子任务 5.3: 根据视频文件名生成对应的txt文件名
            video_path_obj = Path(video_path)
            video_name = video_path_obj.stem  # 获取不带扩展名的文件名
            result_file_path = Path(self.m_resultsDir) / f"{video_name}.txt"
            
            # 确保结果目录存在
            Path(self.m_resultsDir).mkdir(parents=True, exist_ok=True)
            
            # 存储所有帧的结果
            results_lines = []
            
            # 子任务 5.1: 实现逐帧读取循环
            frame_number = 0
            while True:
                ret, frame = cap.read()
                
                # 如果读取失败，说明视频结束
                if not ret:
                    break
                
                frame_number += 1
                
                # 子任务 5.2: 对每一帧调用YOLO模型进行推理
                if self._model is None:
                    raise RuntimeError("YOLO模型未加载，请先调用loadModel()方法")
                
                # 使用YOLO模型进行检测
                detection_results = self._model(frame, verbose=False)
                
                # 子任务 5.2: 提取检测结果中的对象类别和置信度
                detected_objects = []
                for result in detection_results:
                    boxes = result.boxes
                    for box in boxes:
                        # 获取类别ID和置信度
                        class_id = int(box.cls[0])
                        confidence = float(box.conf[0])
                        
                        # 获取类别名称
                        class_name = result.names[class_id]
                        
                        # 格式化为 "对象(置信度)" 格式
                        detected_objects.append(f"{class_name}({confidence:.2f})")
                
                # 子任务 5.2: 格式化为"Frame {帧号}: {对象1}({置信度1}), {对象2}({置信度2})"格式
                if detected_objects:
                    line = f"Frame {frame_number}: {', '.join(detected_objects)}"
                else:
                    line = f"Frame {frame_number}: "
                
                results_lines.append(line)
                
                # 子任务 5.4: 在处理每一帧后调用callback函数（如果提供）
                if callback is not None:
                    callback(frame_number, total_frames)
            
            # 子任务 5.3: 将格式化的结果写入ResultsDirectory，使用UTF-8编码保存文件
            with open(result_file_path, 'w', encoding='utf-8') as f:
                f.write('\n'.join(results_lines))
            
            return True
            
        except Exception as e:
            # 子任务 5.4: 添加视频读取失败、文件写入失败的异常处理
            print(f"处理视频时出错 {video_path}: {str(e)}")
            return False
            
        finally:
            # 子任务 5.4: 确保视频资源正确释放（使用finally块）
            if cap is not None:
                cap.release()
    
    def processVideosFromPath(self, callback: Optional[Callable[[int, int, str], None]] = None) -> dict:
        """
        批量处理videos文件夹中的所有视频
        
        参数:
            callback: 可选的回调函数，签名为 callback(current_video, total_videos, video_name)
        
        返回值:
            dict: 处理结果统计
              {
                  'total': int,      # 总视频数
                  'success': int,    # 成功处理数
                  'failed': list     # 失败的视频列表
              }
        
        功能:
            - 扫描m_videosDir目录
            - 过滤支持的视频格式（.mp4, .avi, .mov, .mkv）
            - 依次调用processVideo处理每个视频
            - 自动创建results目录（如果不存在）
        """
        # 子任务 6.1: 使用pathlib扫描m_videosDir目录
        videos_path = Path(self.m_videosDir)
        
        # 检查视频目录是否存在
        if not videos_path.exists():
            raise FileNotFoundError(f"视频目录不存在: {self.m_videosDir}")
        
        # 子任务 6.1: 过滤支持的视频格式（.mp4, .avi, .mov, .mkv）
        supported_formats = ['.mp4', '.avi', '.mov', '.mkv']
        video_files = []
        for ext in supported_formats:
            video_files.extend(videos_path.glob(f'*{ext}'))
            # 也支持大写扩展名
            video_files.extend(videos_path.glob(f'*{ext.upper()}'))
        
        # 子任务 6.1: 检查并自动创建results目录
        results_path = Path(self.m_resultsDir)
        results_path.mkdir(parents=True, exist_ok=True)
        
        # 子任务 6.2: 统计成功和失败的视频数量
        total_videos = len(video_files)
        success_count = 0
        failed_videos = []
        
        # 子任务 6.2: 遍历所有视频文件，依次调用processVideo
        for index, video_file in enumerate(video_files, start=1):
            video_name = video_file.name
            
            # 子任务 6.2: 在处理每个视频后调用callback函数（如果提供）
            if callback is not None:
                callback(index, total_videos, video_name)
            
            # 调用processVideo处理单个视频
            success = self.processVideo(str(video_file))
            
            if success:
                success_count += 1
            else:
                failed_videos.append(video_name)
        
        # 子任务 6.2: 返回包含total、success、failed的字典
        return {
            'total': total_videos,
            'success': success_count,
            'failed': failed_videos
        }
    
    def searchKeyword(self, keyword: str) -> List[CMatchVideo]:
        """
        在结果文件中搜索关键词
        
        参数:
            keyword: 搜索关键词（不区分大小写）
        
        返回值:
            List[CMatchVideo]: 匹配结果列表，按匹配次数降序排序
        
        功能:
            - 遍历results文件夹中的所有txt文件
            - 搜索包含关键词的行
            - 构建CMatchVideo对象
            - 按匹配次数降序排序
        """
        # 子任务 7.1: 使用pathlib扫描ResultsDirectory中的所有txt文件
        results_path = Path(self.m_resultsDir)
        
        # 检查结果目录是否存在
        if not results_path.exists():
            print(f"结果目录不存在: {self.m_resultsDir}")
            return []
        
        # 获取所有txt文件
        txt_files = list(results_path.glob('*.txt'))
        
        # 存储所有匹配结果
        match_results = []
        
        # 子任务 7.1: 实现不区分大小写的关键词匹配逻辑
        keyword_lower = keyword.lower()
        
        # 遍历每个txt文件
        for txt_file in txt_files:
            try:
                # 子任务 7.2: 对每个匹配的txt文件，收集匹配的行号和内容
                matched_lines = []
                
                # 读取文件内容
                with open(txt_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                
                # 搜索包含关键词的行
                for line_number, line_content in enumerate(lines, start=1):
                    # 不区分大小写的匹配
                    if keyword_lower in line_content.lower():
                        matched_lines.append({
                            'line_number': line_number,
                            'content': line_content.strip()
                        })
                
                # 如果有匹配的行，创建CMatchVideo对象
                if matched_lines:
                    # 子任务 7.2: 根据txt文件名追溯对应的视频文件路径
                    result_file_name = txt_file.name
                    video_base_name = txt_file.stem  # 获取不带扩展名的文件名
                    
                    # 在videos目录中查找对应的视频文件
                    videos_path = Path(self.m_videosDir)
                    video_file = None
                    video_name = None
                    
                    # 尝试查找支持的视频格式
                    supported_formats = ['.mp4', '.avi', '.mov', '.mkv']
                    for ext in supported_formats:
                        potential_video = videos_path / f"{video_base_name}{ext}"
                        if potential_video.exists():
                            video_file = str(potential_video)
                            video_name = potential_video.name
                            break
                        # 也尝试大写扩展名
                        potential_video = videos_path / f"{video_base_name}{ext.upper()}"
                        if potential_video.exists():
                            video_file = str(potential_video)
                            video_name = potential_video.name
                            break
                    
                    # 如果找不到视频文件，使用默认路径
                    if video_file is None:
                        video_file = str(videos_path / f"{video_base_name}.mp4")
                        video_name = f"{video_base_name}.mp4"
                    
                    # 子任务 7.2: 创建CMatchVideo对象，填充所有字段
                    match_video = CMatchVideo(
                        m_videoName=video_name,
                        m_videoPath=video_file,
                        m_resultFile=result_file_name,
                        m_matchedLines=matched_lines,
                        m_matchCnt=len(matched_lines)
                    )
                    
                    match_results.append(match_video)
                    
            except Exception as e:
                print(f"处理结果文件时出错 {txt_file}: {str(e)}")
                continue
        
        # 子任务 7.3: 按m_matchCnt降序排序结果列表
        match_results.sort(key=lambda x: x.m_matchCnt, reverse=True)
        
        # 子任务 7.3: 返回List[CMatchVideo]
        return match_results
